class TransNode:
    """
    the state transition matrix
    """
    def __init__(self,key,slot_type= None, intent_type=None):
        self.key = key
        self.children = {}
        self.slot_type = slot_type
        self.intent_type = intent_type
        self.is_leaf = False
        self.template_id = []
        self.indegree = 0
        self.tp_id_counter = {}

    def copy(self):
        """
        copy itself and create a new node
        :return:
        """
        node = TransNode(self.key,self.slot_type,self.intent_type)
        node.is_leaf = self.is_leaf
        node.indegree = self.indegree
        for tnode in self.children:
            node.children[tnode] = tnode
        for t in self.template_id:
            node.template_id.append(t)
        for t in self.tp_id_counter:
            node.tp_id_counter[t] = self.tp_id_counter
        if hasattr(self,"var_type"):
            node.var_typoe = self.var_type
        if hasattr(self,"pos"):
            node.pos = self.pos
        if hasattr(self,"is_transformer"):
            node.is_transformer = self.is_transformer

    def mount(self,varnode,tp_idx):
        """
        mount the variable entity into the current node
        :param varnode:
        :param tp_idx:
        :return:
        """
        for vnode in varnode.children:
            if vnode not in self.children:
                self.children[vnode] = vnode
            else:
                curvarnode = self.children[vnode]
                if type(curvarnode)==VarNode:
                    curvarnode = curvarnode.to_transnode()
                    curvarnode.indegree = 1
                    curvarnode.tp_id_counter[tp_idx] = 0
                    if curvarnode.pos != vnode.pos:
                        print("error conflicting, two varnode have two different positions")
                del self.children[curvarnode]
                self.children[curvarnode] = curvarnode
            self.merge(curvarnode,vnode,tp_idx)


    def merge(self,transnode,varnode,tp_idx):
        if id(transnode) == id(varnode):
            print("[info] same node",transnode,varnode)
            return
        if type(transnode) == VarNode:
            print("[big bud]")
            return
        if type(varnode) == TransNode:
            print("[bug]")
        if hasattr(transnode,'var_type'):
            pass
        if type(varnode) == VarNode and type(transnode) == TransNode:
            transnode.var_type = varnode.var_type
            for vnode in varnode.children:
                if type(vnode) == TransNode:
                    print("[bigBug] varnode contains transnode")
                    continue
                if vnode not in transnode.children:
                    transnode.children[vnode] = vnode
                else:
                    real_node = transnode.children[vnode]
                    if type(real_node) == VarNode:
                        real_node = real_node.to_transnode()
                        real_node.indegree = 1
                        real_node.tp_id_counter[tp_idx] = 0
                        del transnode.children[real_node]
                        transnode.children[real_node] = real_node
                    self.merge(real_node,vnode,tp_idx)
    def set(self,node):
        if node in self.children:
            del self.children[node]
        self.children[node] = node

    def __hash__(self):
        return hash(self.key)
    def __eq__(self, other):
        if hash(other) == hash(self):
            return True
        return False
    def __str__(self):
        return "{}--{}".format("TransNode",self.key)

class VarNode:
    """
    variable entity, e.g. company name, city name
    """
    def __init__(self,key, slot_type = None, var_type = None):
        self.key = key
        self.children = {}
        self.var_type = var_type
        self.slot_type = slot_type
        self.is_leaf = False
        self.pos = -1
    def __hash__(self):
        return hash(self.key)

    def __eq__(self,other):
        if hash(other) == hash(self):
            return True
        return False
    def __str__(self):
        return "{}-{}".format("varNode",self.key)
    def to_transnode(self):
        node = TransNode(self.key)
        node.children = {}
        for vnode in self.children:
            node.children[vnode] = vnode
        node.slot_type = self.slot_type
        node.is_leaf = self.is_leaf
        node.var_type = self.var_type
        if hasattr(self,"intent_type"):
            node.intent_type = self.intent_type
        if hasattr(self,"template_id"):
            for idx in self.template_id:
                node.template_id.append(idx)
        node.is_transformer = True
        node.pos = self.pos
        node.indegree = 0
        node.tp_id_counter ={}
        return node

    def check(self,varnode):
        for node in varnode.children:
            if type(node) == TransNode:
                return False
            flag = self.check(node)
            if flag == False:
                return False
        return True


class VarTrie:
    """
    the tree for variable entity
    """
    def __init__(self,var_type):
        self.root = VarNode("root",var_type = var_type)
        self.var_type = var_type

    def add(self,words):
        children = self.root.children
        tn = None
        for i,word in enumerate(words):
            if word in children:
                tn = children[word]
                children = tn.children
            else:
                tn = VarNode(word,self.var_type,self.var_type)
                children[tn] = tn
                children = tn.children
        if tn:
            tn.var_type = self.var_type
            tn.is_leaf = True
            tn.pos = len(words)
    def build(self,data):
        for line in data:
            self.add(line)

class DFA:
    def __init__(self,mode = "mix"):
        self.root = TransNode("root")
        self.var_trie = {}
        self.template_num = 0
        self.mode = mode

    def load_var_trie(self,key,var_trie):
        self.var_trie[key] = var_trie
    def add_list(self,prenodes,wlist, tag, tp_idx):
        last_node = []
        for words in wlist:
            if words[0] == "" or words[0] =="#":
                last_node.extend(prenodes)
            elif words =="<*>":
                tn = TransNode(words,tag)
                if hasattr(tn,"low"):
                    print("[error],<*> low high conflict")

                tn.low = wlist[1]
                tn.high = wlist[2]
                postnodes = [tn]
                realnodes = self.set_next_nodes(prenodes,postnodes,tp_idx)
                last_node.extend(realnodes)
                break
            else:
                # error when different type list overlap
                for i,w in enumerate(words):
                    cur = TransNode(w,tag)
                    if i ==0:
                        nest_nodes = self.set_next_nodes(prenodes,[cur],tp_idx)
                    else:
                        nest_nodes = self.set_next_nodes(nest_nodes,[cur],tp_idx)
                last_node.extend(nest_nodes)
        return last_node

    def set_node(self,prenode,postnode,tp_idx):
        if type(prenode) ==VarNode:
            print("[big bug] parent node is varnode")
        if type(postnode) ==VarNode:
            print("[big bug] child node is varnode")

        if tp_idx not in postnode.tp_id_counter:
            postnode.tp_id_counter[tp_idx] = len(postnode.tp_id_counter)
        if postnode in prenode.children:
            tnode = prenode.children[postnode]
            if type(tnode) == VarNode:
                print("[warning] convert varnode into transnode")
                tnode = tnode.to_transnode()
                del prenode.children[tnode]
                prenode.children[tnode] =tnode

            if tp_idx not in tnode.tp_id_counter:
                tnode.tp_id_counter[tp_idx] = len(tnode.tp_id_counter)
            if tnode.slot_type != postnode.slot_type:
                if self.mode == "alone":
                    print("[error]","suffer conflict between two node")
            return tnode
        else:
            prenode.children[postnode] = postnode
            postnode.indegree +=1
            return postnode

    def set_next_nodes(self,prenodes,postnodes,tp_idx):
        real_postnode = []
        visited = {}
        for prenode in prenodes:
            for postnode in postnodes:
                realnode = self.set_node(prenode,postnode,tp_idx)
                if id(realnode) in visited:
                    continue
                else:
                    real_postnode.append(realnode)
                    visited[id(realnode)] = realnode
        return real_postnode

    def mount_varnode(self,prenodes,word,varnode,tp_idx):
        for node in prenodes:
            if word not in node.children:
                if type(node) == VarNode:
                    print("[big error] skip varnod")
                    continue
                node.mount(varnode,tp_idx)
            else:
                pass

    def add(self, words,tags, intent,tp_idx):
        memory = [None]*(len(words)+1)
        prenodes = [self.root]
        memory[0] = prenodes
        for i,word in enumerate(words):
            tag=tags[i]
            if type(word) ==list:
                prenodes = self.add_list(prenodes,word,tag,tp_idx)
            else:
                if word in self.var_trie:
                    self.mount_varnode(prenodes,word,self.var_trie[word],tp_idx)
                tn=TransNode(word,tag)
                postnodes = [tn]
                realnodes = self.set_next_nodes(prenodes,postnodes,tp_idx)
                prenodes = realnodes

            prenodes = self.remove_duplicate_node(prenodes)
            memory[i+1] = prenodes

        for tn in prenodes:
              if type(tn) == VarNode or hasattr(tn,"var_type"):
                  print("[info] set template-idx in varnode")
              if hasattr(tn,"template_id") and len(tn.template_id) != 0:
                  pass

              tn.intent_type = intent
              tn.template_id.append(tp_idx)
              tn.is_leaf =True

    def replace(self,prenodes,node,tp_idx):
        del node.tp_id_counter[tp_idx]
        tnode = node.copy()
        tnode.indegree = 0
        tnode.tp_id_counter={}
        tnode.tp_id_counter[tp_idx] = 0
        for prenode in prenodes:
            if node not in prenode.children:
                continue
            curnode = prenode.children[node]

            if id(curnode) == id(node):
                node.indegree -= 1
                tnode.indegree +=1
                del prenode.children[node]
                prenode.children[tnode] = tnode
        return tnode

    def remove_duplicate_node(self,last_node):
        tmp = []
        visited = []
        for node in last_node:
            if type(node) == VarNode:
                print("[big bug] prenode contains varnode")
            if id(node) in visited:
                continue
            else:
                visited[id(node)] = None
                tmp.append(node)

        if len(last_node) -len(tmp) >0:
            print("[warnings] duplicate num",len(last_node) - len(tmp))
        return tmp

    def build(self,word_seq,tag_seq,intents):
        for i in range(len(word_seq)):
            self.add(word_seq[i],tag_seq[i],intents[i],1)

    def match(self,words):
        """
        match the string with dfa
        :param words:
        :return:
        """
        children = self.root.children
        node = self.root
        tnode = None
        tags=[]
        memory = [node]
        index = 0
        while index < len(words):
            word = words[index]
            index +=1
            if word in children:
                tnode = children[word]
                children = tnode.children
                tags.append(tnode.slot_type)
                memory.append(tnode)
            else:
                if type(tnode) ==VarNode and tnode.is_leaf and tnode.var_type != None:
                    pos = tnode.pos
                    if pos ==-1:
                        return False,tags,None,None
                    node = memory[-pos-1]

                    if node and tnode.var_type in node.children:
                        var_type = tnode.var_type
                        tnode = node.children[tnode.var_type]
                        memory.append(tnode)
                        children = tnode.children
                        varnode_type = tnode.slot_type
                        if word in children:
                            tnode_next = children[word]
                            children = tnode_next.children
                            k = len(tags) -1
                            while pos >0:
                                tags [k-pos+1] = varnode_type
                                pos-=1
                            tags.append(tnode_next.slot_type)
                            tnode = tnode_next
                            memory.append(tnode)
                            continue
                        else:
                            return False,tags,None,None
                    else:
                        return False,tags, None,None
                elif "<*>" in children: # match *
                    tnode = None
                    repeat_num = 0
                    flag = False
                    starnode = children["<*>"]
                    index -=1
                    while repeat_num <= starnode.high and index < len(words):
                        word = words[index]
                        index += 1
                        if starnode.high >= repeat_num >= starnode.low and word in starnode.children:
                            tnode = starnode.children[word]
                            memory.append(tnode)
                            children = tnode.children
                            flag = True
                            tags.append(tnode.slot_type)
                            break
                        tags.append(starnode.slot_type)
                        repeat_num+=1
                    if flag:
                        continue
                return False,tags,None,None
        #step 2 match the last word to check whether it is a node or not

        if type(tnode) ==VarNode:
            if tnode.is_leaf and tnode.var_type != None:
                pos = tnode.pos
                if pos ==-1:
                    return False,tags,None,None
                node = memory[-pos-1]
                if node and tnode.var_type in node.children:
                    var_type = tnode.var_type
                    transnode = node.children[tnode.var_type]
                    node = transnode
                    k = len(tags) -1
                    while pos > 0:
                        tags[k-pos +1] = node.slot_type
                        pos -= 1
                    if node.intent_type is not None:
                        return True, tags, node.intent_type,node.template_id
            return True,tags, None, None
        if tnode.template_id:
            return True,tags, tnode.intent_type, tnode.template_id

        if tnode.is_leaf:
            pos = tnode.pos
            if pos ==-1:
                return False,tags,None,None
            node = memory[-pos-1]
            if node and tnode.var_type in node.children:
                var_type = tnode.var_type
                node = node.children[tnode.var_type]
                k = len(tags)-1
                while pos >0:
                    tags[k-tags+1] =node.slot_type
                    pos -= 1
                if node.intent_type is not None:
                    return  True, tags, node.intent_type,node.template_id
        return False,tags,None,None

if __name__=="__main__":
    template = "<doc_type> this is? [in|from|of] {three copies}:slot_1 ->intent_1"
    company = VarTrie('<company>')
    company_data = [['abc','zhong'],['ccb','ltd']]
    company.build(company_data)
